
from genome.readgenome import readGenomicSequence
import sys
from theconfig import *

def readExonLengths(genefile="../../references/knownGene.table.txt"):
    exonsByEnd = {}
    exonsByStart = {}
    exonBoundaries = {}
    for line in open(genefile):
        if line.startswith("#"): continue
        cols = line.strip().split("\t")
        chr = cols[1]
        strand = cols[2]
        exonCount = int(cols[7])
        exonStarts = cols[8].split(",")[:-1]
        exonEnds = cols[9].split(",")[:-1]
        for i in xrange(exonCount):
            start = long(exonStarts[i])+1
            end = long(exonEnds[i])
            startKey = (chr, start)
            endKey = (chr, end)
            exonsByStart[startKey] = max(exonsByStart.get(startKey, 0), end-start+1)
            exonsByEnd[endKey] = max(exonsByEnd.get(endKey, 0), end-start+1)
            if i>0:
                exonBoundaries[(chr, strand, long(exonEnds[i-1]), long(exonStarts[i])+1)] = True
    return (exonsByStart, exonsByEnd, exonBoundaries)

def generateSpliceJunction(id, leftWidth=None, rightWidth=None, exonsByStart=None, exonsByEnd=None):
    (name, loc) = tuple(id.split(":"))
    (chr, strand, oldLeftWidth, oldRightWidth, leftLoc, rightLoc) = tuple(loc.split("_"))
    leftLoc = long(leftLoc)
    rightLoc = long(rightLoc)
    if leftWidth==None:
        leftWidth = int(oldLeftWidth)
    if rightWidth==None:
        rightWidth = int(oldRightWidth)
    if exonsByEnd!=None:
        if not exonsByEnd.has_key((chr, leftLoc)):
            print >> sys.stderr, "# Warning: no exons ending with ", (chr, leftLoc), "; cf. ", id
            if exonsByEnd.has_key((chr, leftLoc+1)):
                print >> sys.stderr, "### -- Shifting left location by +1 to ", (chr, leftLoc+1)
                leftLoc += 1
            elif exonsByEnd.has_key((chr, leftLoc-1)):
                print >> sys.stderr, "### -- Shifting left location by -1 to ", (chr, leftLoc-1)
                leftLoc -= 1
        if exonsByEnd.has_key((chr, leftLoc)):
            leftWidth = min(leftWidth, exonsByEnd[(chr, leftLoc)])
    if exonsByStart!=None:
        if not exonsByStart.has_key((chr, rightLoc)):
            print >> sys.stderr, "# Warning: no exons starting with ", (chr, rightLoc), "; cf. ", id
            if exonsByStart.has_key((chr, rightLoc+1)):
                print >> sys.stderr, "### -- Shifting right location by +1 to ", (chr, rightLoc+1)
                rightLoc += 1
            elif exonsByStart.has_key((chr, rightLoc-1)):
                print >> sys.stderr, "### -- Shifting right location by -1 to ", (chr, rightLoc-1)
                rightLoc -= 1
        if exonsByStart.has_key((chr, rightLoc)):
            rightWidth = min(rightWidth, exonsByStart[(chr, rightLoc)])
    if leftWidth<=0 or rightWidth<=0:
        print >> sys.stderr, "# Warning: width<=0! ", id
    fullSeq = readGenomicSequence((chr, leftLoc-leftWidth+1, rightLoc+rightWidth-1))
    leftSeq = fullSeq[:leftWidth]
    rightSeq = fullSeq[-rightWidth:]
    spliceSequence = fullSeq[leftWidth:leftWidth+2] + ".." + fullSeq[-rightWidth-2:-rightWidth]
    newId = "%s:%s_%s_%d_%d_%ld_%ld" % (name, chr, strand, leftWidth, rightWidth, leftLoc, rightLoc)
    if (strand=="+" and spliceSequence!="GT..AG") or (strand=="-" and spliceSequence!="CT..AC"):
        print >> sys.stderr, "# Warning: splice sequence ", spliceSequence, " not consistent with strand ", strand, " of ", newId
    newSeq = leftSeq + rightSeq
    return (newId, newSeq)

if __name__=="__main__":
    spliceFile = sys.argv[1]
    leftWidth = rightWidth = long(sys.argv[2])
    debugChanges = (len(sys.argv)>3 and sys.argv[3]=="T")
    (exonsByStart, exonsByEnd, exonBoundaries) = readExonLengths()
    for line in open(spliceFile):
        if line.startswith(">"):
            oldId = line.strip()
            (name, loc) = tuple(oldId.split(":"))
            (chr, strand, oldLeftWidth, oldRightWidth, leftLoc, rightLoc) = tuple(loc.split("_"))            
            key = (chr, strand, long(leftLoc), long(rightLoc))
            if exonBoundaries.has_key(key):
                del exonBoundaries[key]
            else:
                print >> sys.stderr, "# Warning... key", key, "not found in transcripts."
    for key in exonBoundaries.keys():
        print >> sys.stderr, "# Warning... key", key, "not found in splice junction file."
    nn = 0
    for line in open(spliceFile):
        if line.startswith(">"):
            oldId = line.strip()
            (newId, newSeq) = generateSpliceJunction(oldId, leftWidth, rightWidth, exonsByStart, exonsByEnd)
        else:
            nn += 1
            if 0==nn%100:
                print >> sys.stderr, "# Read %d sequences..." % nn
            oldSeq = line.strip()
            if debugChanges and (oldSeq.upper() != newSeq.upper() or oldId != newId):
                print >> sys.stderr, "# Altered splice junction: "
                print >> sys.stderr, oldId
                print >> sys.stderr, oldSeq
                print >> sys.stderr, newId
                print >> sys.stderr, newSeq
            print newId
            print newSeq
